'''
Created on Nov 27, 2011

@author: oabalbin

Returns a pedigree input file for mega2.
pedirgree + marker status.
input pedrigee, pos.base.geno file for each individual in the pedigree

Type    Name    chromosome      hg19.start      hg18.start
M       Marker1 chr1    110201668       110003191


'''
import os
import numpy as np
from collections import defaultdict


class Ped:
    def __init__(self,indiv,markers):
        '''
        '''
        self.indv=indiv
        self.snp_markers=markers
        self.markers_outfile='markers_outfile.txt'
        # create a numpy array for the number of markers.
        self.mat_markers=np.ones((len(indiv), len(markers)*2))*-1
        self.pad=8
        pass
    
    def read_pos(self,file):
        '''
        return a string with all the positions.
        '''
        print file
        ifile = open(file)
        for l in ifile:
            if l.startswith('#'):
                continue
            snps_string=l.strip('\n')
        ifile.close()
        return snps_string
    def find_snp(self,snp, snps_string,pad):
        '''
        '''
        ind=snps_string.find(snp)
        if ind !=-1:
            snp_info = snps_string[ind:ind+len(snp)+pad]
            return snp_info
        else:
            return None
        
    def find_snp_list(self,snpdict,snps_string,individualID):
        '''
        '''
        for marker, snp in snpdict.iteritems():
            snpid=snp['CHRM']+'@'+snp['POS']
            snp_info=self.find_snp(snpid, snps_string,self.pad)
            print marker, snp, snpid
            if snp_info is not None:
                '''
                chr1@868096|T>C|0/1
                '''
                print snp_info
                a1,a2=snp_info.split('|')[-1].split('/')[0],snp_info.split('|')[-1].split('/')[1]
                
                self.mat_markers[individualID,marker*2]=int(a1)+1
                self.mat_markers[individualID,marker*2+1]=int(a2)+1
                #print snp_info, a1,a2,individualID,marker,marker*2, self.mat_markers[individualID,:]
                #print self.mat_markers[individualID,marker:marker+2]
        print self.mat_markers[individualID,:]
    
    def print_mat_markers(self,dir):
        '''
        numpy.savetxt(fname, X, fmt='%.18e', delimiter=' ', newline='\n', header='', footer='', comments='# ')
        Save an array to a text file
        '''
        hd=[]
        for i,marker in self.snp_markers.iteritems():
            a1=marker['type']+str(i)+'.'+marker['type']+'.1'
            a2=marker['type']+str(i)+'.'+marker['type']+'.2'
            hd.append(a1)
            hd.append(a2)
        hd=",".join(hd).replace(',','\t')
        print hd
        of=os.path.join(dir,self.markers_outfile)
        of =open(of,'w')
        of.write(hd+'\n')
        np.savetxt(of,self.mat_markers,fmt='%d',delimiter='\t') #newline='\n',,header=hd
            
    
    def get_markers_status(self,dir):
        '''
        '''
        file_field=2
        for indvID,f in self.indv.iteritems():
            if f[file_field]=='':
                continue
            file=os.path.join(dir,f[file_field])
            snps_string=self.read_pos(file)
            print file
            self.find_snp_list(self.snp_markers,snps_string,indvID)
        
        pass
                
def read_markers(file):
    '''
    '''
    records=defaultdict(dict)
    ifile=open(file)
    ifile.next()
    for i,l in enumerate(ifile):
        f=l.strip('\r\n').split('\t')
        rec={'type':f[0],'Name':f[1],'CHRM':f[2],'POS':f[3],'POSHG18':f[4]}
        records[i]=rec
    ifile.close()
    return records

def read_markers_file(file):
    '''
    '''
    pass
     
def read_patient_files(file):
    ifile=open(file)
    pat=defaultdict()
    i=0
    for l in ifile:
        pat[i]=l.strip('\n').split('\t')
        i+=1
    ifile.close()
    return pat

def main():
    patients_file='/exds/users/oabalbin/projects/fmpn/data/genetic_analysis/patient_pos-base-geno.txt'
    markers_file='/exds/users/oabalbin/projects/fmpn/data/genetic_analysis/fmpn.map.names.txt'
    data_dir='/exds/users/oabalbin/projects/fmpn/data/gatk_calls'
    out_dir='/exds/users/oabalbin/projects/fmpn/data/genetic_analysis'
    indv=read_patient_files(patients_file)
    markers=read_markers(markers_file)
    nindv=len(indv)
    nmarkers=len(markers)
    print indv
    print markers
    fmpns = Ped(indv,markers)
    fmpns.get_markers_status(data_dir)
    fmpns.print_mat_markers(out_dir)
    
main()